#include "detect_nodes.h"
#include "sequence_corpus_reader.h"
#include "gzstream.h"

#include <math.h>
#include <sstream>
#include <boost/lexical_cast.hpp>

int TAG_FREQ_CUTOFF=1;

NodeDetector::NodeDetector(boost::shared_ptr<Parameters> p, bool detect_null)
: myParameters(p), myCorpusReader(new SequenceCorpusReader(p)), 
  myFeatureMap(new FeatureMap())
{
	if (p->feature_map_in_given)
	{
		igzstream fm(p->feature_map_in_arg);
		fm >> *myFeatureMap;
		fm.close();
	}
	if (p->tag_stats_given)
	{
		igzstream in(p->tag_stats_arg);
		myFeatureMap->read_tag_map(in);;
		in.close();
	}

	if (p->label_features_given)
	    read_label_features(p->label_features_arg);
}

NodeDetector::~NodeDetector()
{
	if (myParameters->feature_map_out_given)
	{
		ogzstream fm(myParameters->feature_map_out_arg);
		fm << *myFeatureMap;
		fm.close();
	}
}

void NodeDetector::skip()
{
    SequenceCorpusReader::EntryPtr e = myCorpusReader->next();
}

NodeDetector::ChartPtr NodeDetector::next()
{
	SequenceCorpusReader::EntryPtr e = myCorpusReader->next();
	if (myCorpusReader->eof()) 
		return ChartPtr();

	NodeDetector::ChartPtr chart(new Chart(e->features_chart.size()));
	chart->gold = e->gold;
	chart->order = myCorpusReader->order();
	chart->start_label = e->start_label;
	chart->end_label = e->end_label;
	chart->labels = myCorpusReader->labels();

	int detected_feature = -1;
	size_t num_labels = myCorpusReader->labels();

	// detect features
	for (size_t i=0; i<e->features_chart.size(); ++i)
	{
		for (size_t j=0; j<e->features_chart[i].size(); ++j)
		{
			RowFeaturesPtr rf(new RowFeatures);
			std::set<std::string> &instance = e->features_chart[i][j];

			// find the word in question
			for (std::set<std::string>::iterator it=instance.begin(); it!=instance.end(); ++it)
				if (it->size()>5 && it->substr(0,5)=="WORD=") rf->word=it->substr(5,it->size()-5);

			// get the number of word specific labels, if any
			size_t word_tags = myFeatureMap->word_tags(rf->word);
			if (!word_tags) word_tags = num_labels;
//			std::cerr << rf->word << " " << word_tags << std::endl;

			for (size_t k=0; k<word_tags; ++k)
			{
				NodeFeaturesPtr nf(new NodeFeatures);
				size_t tag = myFeatureMap->map_tag(rf->word,k);
//				std::cerr << "\t" << k << " " << tag << std::endl;
		
				// check if this segment+labelling if gold
				std::pair<size_t,size_t> s_index(i-j,i+1);
//				if(e->gold->count(s_index) > 0 && (*e->gold)[s_index] == k)
				if(e->gold->count(s_index) > 0 && (*e->gold)[s_index] == tag)
				{
					nf->gold = true;
					rf->gold = k;
				}
				else nf->gold = false;

				// map the label to the
				std::string tag_str = boost::lexical_cast<std::string>(tag);

				// detect some node features
				if (myFeatureMap->frozen() || nf->gold)
				{
					for (std::set<std::string>::iterator it=instance.begin();
						it!=instance.end(); ++it)
					{
						detected_feature = myFeatureMap->add
							(std::string(tag_str)+"_"+(*it));
						if (detected_feature >= 0)
							nf->features.insert(std::pair<int,float>
							  (detected_feature,1.0f));
						
						// features conjunctions with sub-label types
//						myLabelFeatures.find(k);
//						if (myLabelFeatures.count(k) > 0)
//						{
//							std::set<std::string> &label_features
//							= myLabelFeatures[k];
//							for (std::set<std::string>::iterator 
//								lf_it=label_features.begin(); 
//								lf_it != label_features.end(); ++lf_it)
//							{
//							detected_feature = myFeatureMap->add
//								(*lf_it+"_"+(*it));
//							if (detected_feature >= 0)
//								nf->features.insert(std::pair<int,float>
//								  (detected_feature,1.0f));
//							}
//						}
					}
				}
				rf->nodes.push_back(nf);
			}
			chart->features[i].push_back(rf);
		}
	}
	return chart;
}

void NodeDetector::detect()
{
}

void NodeDetector::reset()
{
}

void NodeDetector::read_label_features(char *file)
{
    igzstream fd(file);
    int label;
    std::string feature;
    fd >> label;
    while(!fd.eof())
    {
	    fd >> feature;
	    myLabelFeatures[label].insert(feature);
	    fd >> label;
    }
    fd.close();
}

NodeDetector::FeatureSet 
NodeDetector::merge_features(const FeatureSet &lhs, const FeatureSet &rhs)
{
	FeatureSet fs(lhs);
	for(FeatureSet::const_iterator it=rhs.begin(); it != rhs.end(); ++it)
	{
		if(fs.count(it->first) > 0)
			fs[it->first] = fs[it->first] + it->second;	  
		else
			fs[it->first] = it->second;	  
	}
	return fs;
}

void NodeDetector::add_features(FeatureSet &lhs, const FeatureSet &rhs)
{
	for(FeatureSet::const_iterator it=rhs.begin(); it != rhs.end(); ++it)
	{
		if(lhs.count(it->first) > 0)
			lhs[it->first] = lhs[it->first] + it->second;	    
		else
			lhs[it->first] = it->second;	   
	}
}

FeatureMap::FeatureMap()
: myNumFeatures(0), myFrozen(false)
{
}

FeatureMap::~FeatureMap()
{
}

int FeatureMap::add(const std::string &f)
{
	if (myIndex.count(f) == 0)
	{
		if (myFrozen) 
			return -1;
		myIndex[f] = myNumFeatures;
		myReverseIndex[myNumFeatures] = f;
		return myNumFeatures++;
	}
	else
		return myIndex[f]; 
}

int FeatureMap::index(const std::string &f)
{
	if (myIndex.count(f) == 0) return -1;
	return myIndex[f];
}

std::string FeatureMap::string(int i)
{
	if (myReverseIndex.count(i) == 0) return "";
	return myReverseIndex[i];
}

void FeatureMap::read_tag_map(std::istream &in)
{
	int words;
	in >> words;

	for (int i=0; i<words; i++)
	{
		std::string word;
		int freq, tag;
		in >> word >> freq;

		std::string buf;
		std::getline(in, buf);
		std::istringstream ss(buf);

		if (freq > TAG_FREQ_CUTOFF)
		{
			myTagMap[word] = std::vector<int>();
			while(ss >> tag) myTagMap[word].push_back(tag);
		}
	}
}

int FeatureMap::map_tag(const std::string& word, int tag)
{
	if (myTagMap.size()==0 || myTagMap.count(word)==0) 
		return tag;
	
	assert(tag < (int)myTagMap[word].size());
	return myTagMap[word][tag];
}

size_t FeatureMap::word_tags(const std::string& word)
{
	if (myTagMap.size()==0 || myTagMap.count(word)==0) 
		return 0;

	return myTagMap[word].size();
}

std::ostream& operator<<(std::ostream &out, const FeatureMap &fm)
{
	out << fm.myNumFeatures << std::endl;
	for (std::map<std::string,int>::const_iterator it=fm.myIndex.begin();
	     it != fm.myIndex.end(); ++it)
	     out << it->first << ' ' << it->second << std::endl;
	return out;
}

std::istream& operator>>(std::istream &in, FeatureMap &fm)
{
	in >> fm.myNumFeatures;
	std::string key;
	int val;
	in >> key;
	while(!in.eof())
	{
		in >> val;
		fm.myIndex[key] = val;
		in >> key;
	}
	fm.freeze();
	return in;
}
